import logging
from neo4j import GraphDatabase, exceptions
import pandas as pd
import os

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('import_log.log'),
        logging.StreamHandler()
    ]
)


class Neo4jDataImporter:
    def __init__(self, uri, user, password):
        """初始化Neo4j连接"""
        self.uri = uri
        self.user = user
        self.password = password
        self.driver = None

    def connect(self):
        """建立数据库连接并测试连接是否有效"""
        try:
            self.driver = GraphDatabase.driver(self.uri, auth=(self.user, self.password))
            # 测试连接
            with self.driver.session() as session:
                session.run("RETURN 1")
            logging.info("成功连接到Neo4j数据库")
            return True
        except exceptions.AuthError:
            logging.error("认证失败: 用户名或密码不正确")
            logging.error(f"使用的连接信息 - URI: {self.uri}, 用户名: {self.user}")
            return False
        except exceptions.ServiceUnavailable:
            logging.error(f"无法连接到数据库服务，请检查URI: {self.uri}")
            logging.error("确保Neo4j服务已启动并且网络可访问")
            return False
        except Exception as e:
            logging.error(f"连接数据库时发生错误: {str(e)}")
            return False

    def close(self):
        """关闭数据库连接"""
        if self.driver:
            self.driver.close()
            logging.info("数据库连接已关闭")

    def clear_database(self):
        """清空数据库中的所有节点和关系"""
        if not self.driver:
            logging.error("请先建立数据库连接")
            return False

        try:
            with self.driver.session() as session:
                session.run("MATCH (n) DETACH DELETE n")
            logging.info("数据库已清空")
            return True
        except exceptions.Neo4jError as e:
            logging.error(f"清空数据库时发生错误: {str(e)}")
            return False

    def create_constraints(self):
        """创建唯一性约束"""
        if not self.driver:
            logging.error("请先建立数据库连接")
            return False

        try:
            with self.driver.session() as session:
                # 为Person节点的name属性创建唯一性约束
                session.run("CREATE CONSTRAINT IF NOT EXISTS FOR (p:Person) REQUIRE p.name IS UNIQUE")
                # 为Organization节点的name属性创建唯一性约束
                session.run("CREATE CONSTRAINT IF NOT EXISTS FOR (o:Organization) REQUIRE o.name IS UNIQUE")
                # 为Location节点的name属性创建唯一性约束
                session.run("CREATE CONSTRAINT IF NOT EXISTS FOR (l:Location) REQUIRE l.name IS UNIQUE")
            logging.info("已创建唯一性约束")
            return True
        except exceptions.Neo4jError as e:
            logging.error(f"创建约束时发生错误: {str(e)}")
            return False

    def import_data(self, csv_path, has_headers=False):
        """
        从CSV文件导入数据到Neo4j
        :param csv_path: CSV文件路径
        :param has_headers: CSV文件是否包含列名行
        """
        if not self.driver:
            logging.error("请先建立数据库连接")
            return False

        try:
            # 读取CSV文件，根据是否有列名进行处理
            if has_headers:
                df = pd.read_csv(csv_path)
                # 检查是否有至少3列
                if len(df.columns) < 3:
                    logging.error(f"CSV文件需要至少3列数据，实际只有{len(df.columns)}列")
                    return False
            else:
                # 没有列名时，手动指定列名
                df = pd.read_csv(csv_path, header=None, names=['source', 'target', 'relationship'])
                logging.info("CSV文件没有列名，已自动分配列名: source, target, relationship")

            logging.info(f"读取到 {len(df)} 条数据记录")

            # 统计关系类型
            relationship_stats = df['relationship'].value_counts()
            logging.info("关系类型统计:")
            for rel_type, count in relationship_stats.items():
                logging.info(f"  {rel_type}: {count} 条")

            # 导入节点和关系
            batch_size = 1000
            total_imported = 0

            # 创建节点和关系的 Cypher 语句
            create_query = """
            UNWIND $data AS row
            MERGE (s:Person {name: row.source})
            MERGE (t:Person {name: row.target})
            MERGE (s)-[r:RELATIONSHIP {type: row.rel_type}]->(t)
            """

            # 批量导入
            for i in range(0, len(df), batch_size):
                batch = df.iloc[i:i + batch_size]
                data = [
                    {
                        "source": row['source'],
                        "target": row['target'],
                        "rel_type": row['relationship']
                    }
                    for _, row in batch.iterrows()
                ]

                with self.driver.session() as session:
                    session.run(create_query, data=data)

                total_imported += len(batch)
                logging.info(f"已导入 {total_imported}/{len(df)} 条记录")

            logging.info("数据导入完成")

            # 统计导入结果
            with self.driver.session() as session:
                nodes_count = session.run("MATCH (n) RETURN COUNT(n) AS count").single()["count"]
                rels_count = session.run("MATCH ()-[r]->() RETURN COUNT(r) AS count").single()["count"]

            logging.info(f"最终统计: {nodes_count} 个节点, {rels_count} 个关系")
            return True

        except Exception as e:
            logging.error(f"导入过程中发生错误: {str(e)}", exc_info=True)
            return False


if __name__ == "__main__":
    # 直接设置Neo4j连接信息（无需环境变量或dotenv）
    NEO4J_URI = "bolt://localhost:7687"
    NEO4J_USER = "neo4j"
    NEO4J_PASSWORD = "gjj030712"  # 已设置为指定密码
    CSV_PATH = "harrypotterkgdata.csv"  # 数据文件路径
    CSV_HAS_HEADERS = False  # CSV文件没有列名

    # 验证CSV文件是否存在
    if not os.path.exists(CSV_PATH):
        logging.error(f"CSV文件不存在: {CSV_PATH}")
        exit(1)

    # 创建导入器实例并执行导入
    importer = Neo4jDataImporter(NEO4J_URI, NEO4J_USER, NEO4J_PASSWORD)
    try:
        # 建立连接
        if not importer.connect():
            logging.error("无法建立数据库连接，程序退出")
            exit(1)

        # 清空数据库
        if not importer.clear_database():
            exit(1)

        # 创建约束
        if not importer.create_constraints():
            exit(1)

        # 导入数据
        logging.info(f"开始导入数据，文件路径: {CSV_PATH}")
        if not importer.import_data(CSV_PATH, has_headers=CSV_HAS_HEADERS):
            exit(1)

    finally:
        importer.close()